#pragma once

#include "ip_context.hpp"
#include "limbo/customization.hpp"
#include "limbo/details/checksum.h"
#include "limbo/errc.h"
#include "limbo/order.hpp"
#include "limbo/result.h"
#include "opts.h"
#include "proto.h"
#include "state.hpp"
#include <cstdint>

namespace limbo::ip {

template <typename State> struct IpPacket {
  using LowerState = typename State::LowerState;
  using Context = IpContext<State>;
  using Customization = customization<details::base_state_t<State>>;
  static constexpr uint8_t SZ = 20;

  uint8_t ver : 4;
  uint8_t header_length : 4;
  uint8_t tos;
  uint16_t length;

  uint16_t id;
  uint8_t flag : 3;
  uint16_t frag_offset : 13;

  uint8_t ttl;
  uint8_t proto;
  uint16_t checksum;

  IPv4Address source;
  IPv4Address destination;
  Chunk options;
  Chunk payload;

  uint8_t *packet_start; /* needed for icmp */
  LowerState *lower_state;

  template <bool CopyContent = true>
  static Result send(State &state, const Context &ctx, Chunk buff,
                     const Chunk payload) {
    static constexpr uint8_t SZ_WORDS = SZ / sizeof(uint32_t);

    if constexpr (CopyContent) {
      if (buff.size() < payload.size() + SZ) {
        auto code = (uint32_t)Errc::no_enough_space;
        return Result::make_error(code, &state);
      }
    } else if (buff.size() < SZ) {
      auto code = (uint32_t)Errc::no_enough_space;
      return Result::make_error(code, &state);
    }

    auto id = ++state.next_id;

    alignas(4) char data[SZ];
    auto ptr = (unsigned char *)data;
    *ptr++ = (4 << 4) | SZ_WORDS;
    *ptr++ = ctx.tos;

    *((uint16_t *)ptr) = order::native_to_big((uint16_t)(SZ + payload.size()));
    ptr += sizeof(uint16_t);

    *((uint16_t *)ptr) = order::native_to_big(id);
    ptr += sizeof(uint16_t);

    /* flag + frag_offset */
    *((uint16_t *)ptr) = (uint16_t)(ctx.flags << 5);
    ptr += sizeof(uint16_t);

    *ptr++ = ctx.ttl;
    *ptr++ = static_cast<uint8_t>(ctx.proto);

    /* to be calculated later */
    auto checksum = ptr;
    *((uint16_t *)ptr) = 0;
    ptr += 2;

    *((uint32_t *)ptr) = ctx.source.as_net();
    ptr += sizeof(uint32_t);

    *((uint32_t *)ptr) = ctx.destination.as_net();
    ptr += sizeof(uint32_t);

    if constexpr (State::opts & Opts::calculate_checksum) {
      auto value = ~(details::checksum(data, SZ));
      *((uint16_t *)checksum) = order::native_to_big((uint16_t)value);
    } else {
      (void)checksum;
    }

    Customization::copy(buff.data(), data, SZ);
    auto consumed = Chunk(buff.data(), SZ);
    if constexpr (CopyContent) {
      Customization::copy(consumed, payload);
    }
    return Result::make_consumed(&state, consumed);
  }

  Result recv(Chunk buff, State &state, LowerState *lower_state) {
    auto sz = buff.size();
    auto ptr = (unsigned char *)buff.data();
    if (sz < SZ) {
      return Result::make_demand(&state, SZ - sz);
    }
    auto start = ptr;
    ver = (*ptr & (4 << 4)) >> 4;
    if (ver != 4) {
      auto code = (uint32_t)Errc::ip_unsupported_version;
      return Result::make_error(code, &state, Chunk(start, 1));
    }

    header_length = (0b1111) & *ptr;
    if (header_length < 5) {
      auto code = (uint32_t)Errc::ip_header_missized;
      return Result::make_error(code, &state, Chunk(start, 1));
    }
    ++ptr;

    auto header_sz = header_length * sizeof(uint32_t);
    if (header_sz > sz) {
      return Result::make_demand(&state, header_sz - sz);
    }

    tos = *ptr++;

    uint16_t val_16;
    Customization::copy(&val_16, ptr, sizeof(val_16));
    ptr += sizeof(val_16);
    length = order::big_to_native(val_16);
    if (length > sz) {
      return Result::make_demand(&state, length - sz);
    }

    Customization::copy(&val_16, ptr, sizeof(val_16));
    ptr += sizeof(val_16);
    id = order::big_to_native(val_16);

    Customization::copy(&val_16, ptr, sizeof(val_16));
    ptr += sizeof(val_16);
    val_16 = order::big_to_native(val_16);
    /*  */ flag = (val_16 & 0b1110000000000000) >> 13;
    frag_offset = (val_16 & 0b0001111111111111);

    ttl = *ptr++;
    proto = *ptr++;

    Customization::copy(&val_16, ptr, sizeof(val_16));
    ptr += sizeof(val_16);
    checksum = order::big_to_native(val_16);

    uint32_t val_32;
    Customization::copy(&val_32, ptr, sizeof(val_32));
    ptr += sizeof(val_32);
    source = IPv4Address::from_network(val_32);

    Customization::copy(&val_32, ptr, sizeof(val_32));
    ptr += sizeof(val_32);
    destination = IPv4Address::from_network(val_32);

    options = Chunk(start + SZ, header_sz - SZ);
    ptr = options.data() + options.size();

    payload = Chunk(ptr, length - header_sz);
    auto packet = Chunk(start, length);

    if constexpr (State::opts & Opts::validate_checksum) {
      auto cs_raw = details::checksum(start, header_sz);
      auto cs_diff = details::checksum_sub(cs_raw, checksum);
      auto calc_cs = (uint16_t)~cs_diff;
      if (calc_cs != checksum) {
        auto code = (uint32_t)Errc::ip_checksum_mismatch;
        return Result::make_error(code, &state, packet);
      }
    }

    /* success */
    this->lower_state = lower_state;
    packet_start = start;
    return Result::make_consumed(&state, packet);
  }
};

template <typename State>
inline static uint32_t get_prefix(const IpContext<State> &) {
  return IpPacket<State>::SZ;
}

template <typename State>
inline static Chunk get_payload(IpPacket<State> &packet) {
  return packet.payload;
}

} // namespace limbo::ip
